import numpy as np 
import cv2, matplotlib
from pathlib import Path
import imageio.v3 as iio 
import matplotlib.pyplot as plt
import laserbeamsize as lbs
from tkinter.filedialog import askopenfilenames
from scipy.optimize import curve_fit

def pseudo_voigt(x, *params):
    A, mu, sigma, eta, asy_ratio, offset = params
    sigma += asy_ratio * (x - mu)
    gamma = sigma * (2 * np.log(2))**0.5
    gaussian = np.exp(-(x-mu)**2/(2*sigma**2)) / (sigma * (2*np.pi)**0.5)
    lorentzian = gamma / (np.pi * ((x-mu)**2 + gamma**2))
    return A * (eta * gaussian + (1 - eta)*lorentzian) + offset

def find_spot(img, thresh=None):
    if ~thresh:
        thresh = np.mean(img)*1.5
    _, mask = cv2.threshold(img, thresh, 255, cv2.THRESH_BINARY)
    mask = mask.astype('uint8')
    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    areas = [cv2.contourArea(i) for i in contours]
    max_id = np.argmax(areas)
    max_contour = contours[max_id].squeeze()
    center = np.mean(max_contour, 0)
    return [int(i) for i in center]


def main(path, crop_size=256, pixel_size=3.25, fit_projection=True, thresh=None, re_crop=False):
    img = iio.imread(str(path)).squeeze().astype('uint16')
    img = img[150:-150]
    img = cv2.medianBlur(img, 3)
    center = find_spot(img, thresh)
    img = img[center[1]-(crop_size>>1):center[1]+(crop_size>>1),
              center[0]-(crop_size>>1):center[0]+(crop_size>>1)]
    
    img = lbs.subtract_iso_background(img)

    fig, ax = plt.subplots(1, 2, figsize=[12, 5])
    fig.subplots_adjust(left=0.05, right=0.95)
    fig.suptitle(f'Spot Analysis ({path.name})', fontsize=16)
    ax[0].set(title='Image')
    ax[0].axis('off')
    im = ax[0].imshow(img)
    cbar = fig.colorbar(im, label='intensity')
    cbar.formatter.set_powerlimits((-3, 4))

    if re_crop:
        img = lbs.crop_image_to_integration_rect(img, *lbs.beam_size(img))[0]
    v, h = img.shape
    xs_v = (np.arange(v) - v/2) * pixel_size
    xs_h = (np.arange(h) - h/2) * pixel_size

    if fit_projection:
        ys_h = np.sum(img, axis=0)
        ys_v = np.sum(img, axis=1)
    else:
        y, x = np.argwhere(img == np.max(img))[0]
        ys_h = np.mean(img[y-1:y+2, :], axis=0)
        ys_v = np.mean(img[:, x-1:x+2], axis=0)

    for xs, ys, tag in zip([xs_h, xs_v], [ys_h, ys_v], ['_h', '_v']):
        popt, _ = curve_fit(pseudo_voigt, xs, ys,
                            p0=[np.sum(ys*pixel_size), np.mean(xs), np.ptp(xs)/10, 0.5, 0, 0],
                            bounds=([0, -np.inf, 0, 0, -np.inf, -np.inf],
                                    [np.inf, np.inf, np.inf, 1, np.inf, np.inf]),
                            )
        ax[1].scatter(xs, ys, s=10, label=f'FWHM{tag}: {popt[2]*2.355:.2f} $\\mu m$')
        ax[1].plot(xs, pseudo_voigt(xs, *popt))
    
    ax[1].set(title='Profile', xlabel='$\\mu m$', ylabel='Intensity')
    ax[1].set_title(f' Flux: {np.sum(img):.2e}', y=0.9, loc='left')
    ax[1].ticklabel_format(axis='y', style='sci', scilimits=(-3, 4))
    ax[1].legend()
    if fit_projection:
        ax[1].set_title(f'Fit projection ', y=0.1, loc='right')
        fig.savefig(str(path.with_name(path.stem + '_fit projection.png')))
    else:
        ax[1].set_title(f'Fit center line ', y=0.1, loc='right')
        fig.savefig(str(path.with_name(path.stem + '_fit fit center line.png')))   
    plt.show()
    

if __name__ == '__main__':
    files = askopenfilenames(filetypes=[('tif', ['tif', 'tiff'])])
    if len(files) > 1:
        matplotlib.use('agg')
    for file in files:
        try:
            main(Path(file), crop_size=150, fit_projection=1, thresh=0)
        except Exception as e:
            print(e)